#include "MatchShapeBased.h"

#include <QDebug>
#include <random>
#include <QDir>
#include "Performance.h"
#include "VisionApi.h"

MatchShapeBased::MatchShapeBased()
{
}

MatchShapeBased::~MatchShapeBased()
{
}

int MatchShapeBased::execute(cv::Mat src)
{
    Performance timeTest;
    timeTest.Reset();

    allResults_.clear();

    if (isDrawResult())
    {
         src.copyTo(resultImage_);
    }

    HImage srcImg = VisionApi::HImageFromMat(src);

    srcImg = srcImg.MedianImage("circle", 1.5, "mirrored");
    matchRoi_.GenRectangle1(searchRoi_.tl().y, searchRoi_.tl().x, searchRoi_.br().y, searchRoi_.br().x);
    HImage roiImg = srcImg.ReduceDomain(matchRoi_);
    

    HTuple modelIds;
    HTuple r, c, a, s, score;
    roiImg.FindScaledShapeModels(multShapeModels_, HTuple(1).TupleRad(), HTuple(360).TupleRad().D(),
        0.8, 1.2, 0.55, 0, 0.5, HTuple("least_squares").S(),
        0, 0.9, &r, &c, &a, &s, &score, &modelIds);

    // 记录匹配物料数目
    matchCount_ = r.Length();
    for (int i = 0; i < matchCount_; i++)
    {
        MatchResult mr;
        mr.index = i;
        mr.x = c[i].D();
        mr.y = r[i].D();
        mr.score = score[i].D();
        mr.rotAngle = a[i].D();
        mr.scale = s[i].D();
        mr.modelId = modelIds[i].I();

        allResults_.push_back(mr);
    }

    selectResult_ = findBestResult();   

    /*if (selectResult_)
    {
        calibration_->affineTransToRobot(bestResult_->x, bestResult_->y);
        findObj_ = true;
    }*/

    
    // 绘制最佳结果（分数最高）
    drawMatchResult(src, true);   

    // 将目标旋转角度统一到：[-90° 90°]
    for (auto& m : allResults_)
    {
        m.rotAngle = VisionApi::rad2Deg(m.rotAngle);
        if (m.rotAngle > 270)
        {
            m.rotAngle -= 360;
        }
        else if (m.rotAngle > 90)
        {
            m.rotAngle -= 180;
        }
        m.rotAngle *= -1;
    }

    inspectTime_ = timeTest.GetTimeMilli();
    return 0;
}

int MatchShapeBased::createScaledShapeModel(cv::Mat src, cv::Rect roi)
{
    HImage modelImg = VisionApi::HImageFromMat(src);
    modelImg = modelImg.MedianImage("square", 1.0, "mirrored");

    HXLDCont modelContour;
    std::vector<HShapeModel> shapeModels;
    allModelContours_.clear();

    HRegion roiRegion;
    roiRegion.GenRectangle1(roi.y, roi.x, roi.br().y, roi.br().x);

    auto reduce1 = modelImg.ReduceDomain(roiRegion);

    HShapeModel sh1 = reduce1.CreateScaledShapeModel(HTuple("auto"), HTuple(1).TupleRad(), HTuple(360).TupleRad(), HTuple("auto"),
        HTuple(0.8), HTuple(1.2), HTuple("auto"), HTuple("auto"), HTuple("use_polarity"), HTuple("auto"), HTuple(10));
    shapeModels.push_back(sh1);

    modelContour = sh1.GetShapeModelContours(1);
    allModelContours_.push_back(modelContour);
    

    int modelCount = shapeModels.size();
    auto p = new HShapeModel[modelCount];
    for (int i = 0; i < modelCount; i++)
    {
        p[i] = shapeModels[i];
    }
    multShapeModels_.Clear();
    // 建立多模板
    multShapeModels_ = HShapeModelArray(p, modelCount);

    return 0;
}

bool MatchShapeBased::loadModel(std::string modelPath)
{
    if (!VisionApi::isFileExist(modelPath.c_str()))
    {
        qWarning() << "模板文件：" << modelPath << "-不存在，请检查路径是否正确";
        return false;
    }

    HTuple modelId;
    HalconCpp::ReadShapeModel(modelPath.c_str(), &modelId);
    if (modelId.Length() == 0)
    {
        qWarning() << "模板加载失败";
        return false;
    }
    else
    {
        qInfo() << modelPath + ":加载成功";
    }

    multShapeModels_.SetFromTuple(modelId);

    return true;
}

bool MatchShapeBased::loadModels(QString modelDir)
{
    QStringList allShapeModelFile;
    QDir dir(modelDir); // 替换为你的目录路径
    QStringList files = dir.entryList(QDir::Files); // 只获取文件列表
    foreach(const QString & file, files) {
        if (file.right(4) == ".shm")
        {
            allShapeModelFile.append(file);
        }
    }

    if (allShapeModelFile.size() < 1)
    {
        qInfo() << "模板文件加载失败，请检查目录文件是否存在";
        return false;
    }
    multShapeModels_.Clear();
    HTuple allModelId;
    for (auto& sh : allShapeModelFile)
    {
        QString modelPath = modelDir + "//" + sh;

        HTuple modelId;
        HalconCpp::ReadShapeModel(modelPath.toStdString().c_str(), &modelId);
        if (modelId.Length() == 0)
        {
            qInfo() << "模板加载失败";
            return false;
        }
        else
        {
            qInfo() << sh + "-加载成功";
        }

        allModelId.Append(modelId);
    }
    multShapeModels_.SetFromTuple(allModelId);

    return true;
}

void MatchShapeBased::saveModel(const char* modelPath)
{
    int modelCount = multShapeModels_.Length();

    for (int i = 0; i < modelCount; i++)
    {
        QString shapeModelPath = modelPath + QString("-") + QString::number(i) + ".shm";
        HalconCpp::WriteShapeModel(multShapeModels_.ConvertToTuple()[i], shapeModelPath.toStdString().c_str());
    }
}

MatchResult* MatchShapeBased::findBestResult()
{
    MatchResult* r = nullptr;
    if (allResults_.size() == 1)
    {
        r = &allResults_[0];
    }
    else if (allResults_.size() > 1)
    {
        r = &allResults_[0];
        for (auto& m : allResults_)
        {
            if (m.x < r->x)
            {
                r = &m;
            }
        }
    }
    return r;
}

void MatchShapeBased::findObjects(HImage src)
{
    // 验证模板
    HTuple modelIds;
    HTuple r, c, a, s, score;
    src.FindScaledShapeModels(multShapeModels_, HTuple(1).TupleRad(), HTuple(360).TupleRad().D(),
        0.8, 1.2, 0.85, 0, 0.5, HTuple("least_squares").S(),
        0, 0.7, &r, &c, &a, &s, &score, &modelIds);


    cv::Mat destImg /*= src.clone()*/;
    cv::cvtColor(destImg, destImg, cv::COLOR_GRAY2BGR);
    int idCount = modelIds.Length();
    for (int i = 0; i < idCount; i++)
    {
        int id = modelIds[i];
        HXLDCont transXld;
        HHomMat2D homMat;
        homMat.HomMat2dIdentity();
        homMat = homMat.HomMat2dTranslate(r[i].D(), c[i].D());
        homMat = homMat.HomMat2dRotate(a[i].D(), r[i].D(), c[i].D());
        homMat = homMat.HomMat2dScale(s[i].D(), s[i].D(), r[i].D(), c[i].D());

        HalconCpp::AffineTransContourXld(allModelContours_[id], &transXld, homMat);
        VisionApi::drawHalconXld(destImg, transXld);
        //drawResult(destImg, transXld);
    }
}

void MatchShapeBased::drawMatchResult(cv::Mat src, bool srcChanged)
{
    if (selectResult_)
    {
        int i = selectResult_->index;
        std::vector<HXLDCont> allContours;
        HXLDCont modelContours;
        HalconCpp::GetShapeModelContours(&modelContours, multShapeModels_.ConvertToTuple(), 1);
        allContours.push_back(modelContours);

       
        if (srcChanged)
        {
            resultImage_ = src;
        }
        else
        {
            src.copyTo(resultImage_);    
        }        
        cv::cvtColor(resultImage_, resultImage_, cv::COLOR_GRAY2BGR);

        // 计算变换矩阵
        int id = selectResult_->modelId;
        HXLDCont transXld;
        HHomMat2D homMat;
        homMat.HomMat2dIdentity();
        homMat = homMat.HomMat2dTranslate(selectResult_->y, selectResult_->x);
        homMat = homMat.HomMat2dRotate(selectResult_->rotAngle, selectResult_->y, selectResult_->x);
        homMat = homMat.HomMat2dScale(selectResult_->scale, selectResult_->scale, selectResult_->y, selectResult_->x);

        // 将halcon轮廓转为opencv轮廓点
        HalconCpp::AffineTransContourXld(allContours[id], &transXld, homMat);

        // 绘制模板轮廓
        // 随机取色
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<int> dis(100, 255);
        auto color = cv::Scalar(dis(gen),
            dis(gen),
            dis(gen));
        // 绘制
        VisionApi::drawHalconXld(resultImage_, transXld, color);

        // 绘制中心点
        double r1 = 0, c1 = 0;
        HTuple outRow, outCol;
        HalconCpp::AffineTransPixel(homMat, r1, c1, &outRow, &outCol);
        cv::Point pt(outCol.D(), outRow.D());
        VisionApi::drawCross(resultImage_, pt);
    }
}

